#!/usr/bin/python
# -*- coding: utf-8 -*-
from typing import List

import pandas as pd

from core.dataClasses import StockTradeDataColumnName
from infrastructure.util import dateutils
from core.tactic import Period, Tactic, AbstractDelayTactic
from index.ma import SMA, SMA_COLUMN_PREFIX
from core.inventedtrader import TradeType


class SMATactic(Tactic):
    """
    简单随机平均线策略：
        买入：给定均价线上涨，收盘价突破给定均价线
        卖出：给定均价线下跌，收盘价跌下给定均价线
    """

    def __init__(self):
        super(SMATactic, self).__init__()

    def setup(self, context):
        df = context.data  # type:pd.DataFrame
        column = SMA_COLUMN_PREFIX + str(self.period)
        if column not in df.columns.values:
            SMA(self.period, key=StockTradeDataColumnName.CLOSE).compute(context.data)
        if self.order_type != TradeType.SELL:
            self.buy_date = df[
                (df[column] > df[column].shift(1))  # 均价线上涨
                & (df[StockTradeDataColumnName.LOW].shift(1) < df[column].shift(1)) & (df[StockTradeDataColumnName.CLOSE] > df[column])  # 基于收盘价判断估价是否上穿均价线
                ].index.strftime(dateutils.DateFormat.Y_m_d).values
        if self.order_type != TradeType.BUY:
            self.sell_date = df[
                (df[column] < df[column].shift(1))  # 均价线下跌
                & (df[StockTradeDataColumnName.HIGH].shift(1) >= df[column].shift(1)) & (df[StockTradeDataColumnName.CLOSE] < df[column])  # 基于收盘价判断估价是否下穿均价线
                ].index.strftime(dateutils.DateFormat.Y_m_d).values


class SMACrossTactic(Tactic):
    """
    简单随机平均线金叉死叉策略
    """

    def __init__(self, period1, period2):
        super(SMACrossTactic, self).__init__()
        self.period1 = period1
        self.period2 = period2

    def setup(self, context):
        df = context.data
        assert (context.data is not None)
        if SMA_COLUMN_PREFIX + str(self.period1) not in df.columns.values:
            SMA(self.period1, key=StockTradeDataColumnName.CLOSE).compute(df)
        if SMA_COLUMN_PREFIX + str(self.period2) not in df.columns.values:
            SMA(self.period2, key=StockTradeDataColumnName.CLOSE).compute(df)

        cross_data1 = df[SMA_COLUMN_PREFIX + str(self.period1)] < df[
            SMA_COLUMN_PREFIX + str(self.period2)]  # type:DataFrame
        cross_data2 = df[SMA_COLUMN_PREFIX + str(self.period1)] >= df[
            SMA_COLUMN_PREFIX + str(self.period2)]  # type:DataFrame
        if self.order_type != TradeType.SELL:
            self.buy_date = df[cross_data2 & cross_data1.shift(-1)].index.strftime(dateutils.DateFormat.Y_m_d).values
        if self.order_type != TradeType.BUY:
            self.sell_date = df[cross_data1 & cross_data2.shift(-1)].index.strftime(dateutils.DateFormat.Y_m_d).values


class SMACrossDelayTactic(AbstractDelayTactic):
    """
    简单随机平均线金叉死叉，延迟策略
    """

    def __init__(self, period1, period2):
        super(SMACrossDelayTactic, self).__init__()
        self.period1 = period1
        self.period2 = period2

    def setup(self, context):
        df = context.data
        assert (context.data is not None)
        if SMA_COLUMN_PREFIX + str(self.period1) not in df.columns.values:
            SMA(self.period1, key=StockTradeDataColumnName.CLOSE).compute(df)
        if SMA_COLUMN_PREFIX + str(self.period2) not in df.columns.values:
            SMA(self.period2, key=StockTradeDataColumnName.CLOSE).compute(df)

        cross_data1 = df[SMA_COLUMN_PREFIX + str(self.period1)] < df[
            SMA_COLUMN_PREFIX + str(self.period2)]  # type:DataFrame
        cross_data2 = df[SMA_COLUMN_PREFIX + str(self.period1)] >= df[
            SMA_COLUMN_PREFIX + str(self.period2)]  # type:DataFrame
        if self.order_type != TradeType.SELL:
            self.buy_date = df[cross_data2 & cross_data1.shift(-1)].index.strftime(dateutils.DateFormat.Y_m_d).values
        if self.order_type != TradeType.BUY:
            self.sell_date = df[cross_data1 & cross_data2.shift(-1)].index.strftime(
                dateutils.DateFormat.Y_m_d).values

